"""Utility helpers for image encoding and related small helpers.

Contains `encode_png` used by the client node to prepare images for ZMQ.
"""
from __future__ import annotations
import numpy as np
import cv2


def encode_png(img: np.ndarray, is_depth: bool = False) -> bytes | None:
    """Encode a numpy image to PNG bytes.

    - Color images expected in RGB (H,W,3) and are converted to BGR for OpenCV.
    - Depth images: floats are scaled to millimeters and converted to uint16.

    Returns bytes on success, or None on failure.
    """
    if is_depth:
        if img.dtype == np.float32 or img.dtype == np.float64:
            arr = (img * 1000.0).astype(np.uint16)
        elif img.dtype == np.uint16:
            arr = img
        else:
            arr = (img.astype(np.float32) * 1000.0).astype(np.uint16)
        ok, buf = cv2.imencode('.png', arr)
    else:
        # assume RGB input, convert to BGR for cv2
        try:
            bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        except Exception:
            # fall back: if single-channel or already BGR, try direct encode
            try:
                ok, buf = cv2.imencode('.png', img)
            except Exception:
                return None
            if not ok:
                return None
            return buf.tobytes()
        ok, buf = cv2.imencode('.png', bgr)

    if not ok:
        return None
    return buf.tobytes()
